#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy as np
import time
from gnuradio import gr, uhd, blocks


class USRPSyncTest(gr.top_block):
    def __init__(self):
        gr.top_block.__init__(self, "USRP时钟同步测试")

        # 硬件参数
        self.samp_rate = 1e6
        self.center_freq = 915e6
        self.tx_delay = 0.1  # 100ms后发送
        self.trigger_gap = 0.005  # 5ms触发间隔

        # 初始化USRP
        self.tx_usrp = uhd.usrp_sink(
            device_addr="serial=31A92E2",
            stream_args=uhd.stream_args(cpu_format="fc32"),
            antenna="TX/RX"
        )
        self.rx_usrp = uhd.usrp_source(
            device_addr="serial=31AC8A5",
            stream_args=uhd.stream_args(cpu_format="fc32"),
            antenna="RX2"
        )

        # 配置通用参数
        for usrp in [self.tx_usrp, self.rx_usrp]:
            usrp.set_samp_rate(self.samp_rate)
            usrp.set_center_freq(uhd.tune_request(self.center_freq))
            usrp.set_gain(20)

        # 强制时钟同步（关键步骤！）
        self._sync_clocks()

        # 流图组件
        self.tx_signal = np.exp(1j * np.linspace(0, 100, 1000)).astype(np.complex64)
        self.tx_src = blocks.vector_source_c(self.tx_signal, False)
        self.rx_sink = blocks.vector_sink_c()

        # 连接流图
        self.connect(self.tx_src, self.tx_usrp)
        self.connect(self.rx_usrp, self.rx_sink)

    def _sync_clocks(self):
        """使用内部时钟的简化同步"""
        print("=== 使用内部时钟同步 ===")

        # 统一设置内部源
        for usrp in [self.tx_usrp, self.rx_usrp]:
            usrp.set_clock_source("internal")
            usrp.set_time_source("internal")
            usrp.set_time_now(uhd.time_spec(0.0))

        # 简单延迟等待
        time.sleep(0.1)

        # 验证
        tx_time = self.tx_usrp.get_time_now().get_real_secs()
        rx_time = self.rx_usrp.get_time_now().get_real_secs()
        self.tx_rx_time_offset = tx_time-rx_time
        print(f"内部时钟同步结果 - 差值: {self.tx_rx_time_offset}s")



    def run_test(self):
        """执行同步测试"""
        print("\n=== 启动同步测试 ===")

        # 设置精确发射时间（当前时间+100ms）
        tx_time = self.tx_usrp.get_time_now() + uhd.time_spec(self.tx_delay)
        # rx_time = tx_time + uhd.time_spec(self.tx_rx_time_offset)  # 补偿发送接收的时钟偏移
        rx_time = self.rx_usrp.get_time_now() + uhd.time_spec(self.tx_delay)  # 补偿发送接收的时钟偏移
        self.tx_usrp.set_start_time(tx_time)
        self.rx_usrp.set_start_time(rx_time)  # 接收同步

        print(f"计划发射时间: {tx_time.get_real_secs():.6f}s")

        # 启动流图
        self.start()
        start_wall_time = time.time()

        # 等待测试完成
        while (time.time() - start_wall_time) < (self.tx_delay + 0.1):
            tx_now = self.tx_usrp.get_time_now().get_real_secs()
            rx_now = self.rx_usrp.get_time_now().get_real_secs()
            print(f"硬件时钟 - TX: {tx_now:.6f}s, RX: {rx_now:.6f}s")
            time.sleep(0.02)

        self.stop()
        self.wait()

        # 分析结果
        rx_data = np.array(self.rx_sink.data())
        if len(rx_data) > 0:
            print(f"\n测试结果:")
            print(f"理论首采样时间: {tx_time.get_real_secs():.6f}s")
        else:
            print("错误: 未接收到数据！")


if __name__ == "__main__":
    tb = USRPSyncTest()
    tb.run_test()